#include "types.h"

#include <algorithm>

#include "console.h"
#include "filehelper.h"
#include "jansson/jansson_x64dbg.h"
#include "stringutils.h"
#include "threading.h"

using namespace Types;

static TypeManager typeManager;

TypeManager::TypeManager() {
  auto p = [this](const std::string& n, Primitive p, int size) {
    primitivesizes[p] = size;
    auto splits = StringUtils::Split(n, ',');
    for (const auto& split : splits) addType("", p, split);
  };
  p("int8_t,int8,char,byte,bool,signed char", Int8, sizeof(char));
  p("uint8_t,uint8,uchar,unsigned char,ubyte", Uint8, sizeof(unsigned char));
  p("int16_t,int16,wchar_t,char16_t,short", Int16, sizeof(short));
  p("uint16_t,uint16,ushort,unsigned short", Int16, sizeof(unsigned short));
  p("int32_t,int32,int,long", Int32, sizeof(int));
  p("uint32_t,uint32,unsigned int,unsigned long", Uint32, sizeof(unsigned int));
  p("int64_t,int64,long long", Int64, sizeof(long long));
  p("uint64_t,uint64,unsigned long long", Uint64, sizeof(unsigned long long));
  p("dsint", Dsint, sizeof(void*));
  p("duint,size_t", Duint, sizeof(void*));
  p("float", Float, sizeof(float));
  p("double", Double, sizeof(double));
  p("ptr,void*", Pointer, sizeof(void*));
  p("char*,const char*", PtrString, sizeof(char*));
  p("wchar_t*,const wchar_t*", PtrWString, sizeof(wchar_t*));
}

bool TypeManager::AddType(const std::string& owner, const std::string& type,
                          const std::string& name) {
  if (owner.empty()) return false;
  validPtr(type);
  auto found = types.find(type);
  if (found == types.end()) return false;
  return addType(owner, found->second.primitive, name);
}

bool TypeManager::AddStruct(const std::string& owner, const std::string& name) {
  StructUnion s;
  s.name = name;
  s.owner = owner;
  return addStructUnion(s);
}

bool TypeManager::AddUnion(const std::string& owner, const std::string& name) {
  StructUnion u;
  u.owner = owner;
  u.name = name;
  u.isunion = true;
  return addStructUnion(u);
}

bool TypeManager::AddMember(const std::string& parent, const std::string& type,
                            const std::string& name, int arrsize, int offset) {
  if (!isDefined(type) && !validPtr(type)) return false;
  auto found = structs.find(parent);
  if (arrsize < 0 || found == structs.end() || !isDefined(type) ||
      name.empty() || type.empty() || type == parent)
    return false;
  auto& s = found->second;

  for (const auto& member : s.members)
    if (member.name == name) return false;

  auto typeSize = Sizeof(type);
  if (arrsize) typeSize *= arrsize;

  Member m;
  m.name = name;
  m.arrsize = arrsize;
  m.type = type;
  m.offset = offset;

  if (offset >= 0)  // user-defined offset
  {
    if (offset < s.size) return false;
    if (offset > s.size) {
      Member pad;
      pad.type = "char";
      pad.arrsize = offset - s.size;
      char padname[32] = "";
      sprintf_s(padname, "padding%d", pad.arrsize);
      pad.name = padname;
      s.members.push_back(pad);
      s.size += pad.arrsize;
    }
  }

  s.members.push_back(m);

  if (s.isunion) {
    if (typeSize > s.size) s.size = typeSize;
  } else {
    s.size += typeSize;
  }
  return true;
}

bool TypeManager::AppendMember(const std::string& type, const std::string& name,
                               int arrsize, int offset) {
  return AddMember(laststruct, type, name, arrsize, offset);
}

bool TypeManager::AddFunction(const std::string& owner, const std::string& name,
                              const std::string& rettype,
                              CallingConvention callconv, bool noreturn) {
  auto found = functions.find(name);
  if (found != functions.end() || name.empty() || owner.empty()) return false;
  lastfunction = name;
  Function f;
  f.owner = owner;
  f.name = name;
  if (rettype != "void" && !isDefined(rettype) && !validPtr(rettype))
    return false;
  f.rettype = rettype;
  f.callconv = callconv;
  f.noreturn = noreturn;
  functions.insert({f.name, f});
  return true;
}

bool TypeManager::AddArg(const std::string& function, const std::string& type,
                         const std::string& name) {
  if (!isDefined(type) && !validPtr(type)) return false;
  auto found = functions.find(function);
  if (found == functions.end() || function.empty() || name.empty() ||
      !isDefined(type))
    return false;
  lastfunction = function;
  Member arg;
  arg.name = name;
  arg.type = type;
  found->second.args.push_back(arg);
  return true;
}

bool TypeManager::AppendArg(const std::string& type, const std::string& name) {
  return AddArg(lastfunction, type, name);
}

int TypeManager::Sizeof(const std::string& type) const {
  auto foundT = types.find(type);
  if (foundT != types.end()) return foundT->second.size;
  auto foundS = structs.find(type);
  if (foundS != structs.end()) return foundS->second.size;
  auto foundF = functions.find(type);
  if (foundF != functions.end()) {
    const auto foundP = primitivesizes.find(Pointer);
    if (foundP != primitivesizes.end()) return foundP->second;
    return sizeof(void*);
  }
  return 0;
}

bool TypeManager::Visit(const std::string& type, const std::string& name,
                        Visitor& visitor) const {
  Member m;
  m.name = name;
  m.type = type;
  return visitMember(m, visitor);
}

template <typename K, typename V>
static void filterOwnerMap(std::unordered_map<K, V>& map,
                           const std::string& owner) {
  for (auto i = map.begin(); i != map.end();) {
    auto j = i++;
    if (j->second.owner.empty()) continue;
    if (owner.empty() || j->second.owner == owner) map.erase(j);
  }
}

void TypeManager::Clear(const std::string& owner) {
  laststruct.clear();
  lastfunction.clear();
  filterOwnerMap(types, owner);
  filterOwnerMap(structs, owner);
  filterOwnerMap(functions, owner);
}

template <typename K, typename V>
static bool removeType(std::unordered_map<K, V>& map, const std::string& type) {
  auto found = map.find(type);
  if (found == map.end()) return false;
  if (found->second.owner.empty()) return false;
  map.erase(found);
  return true;
}

bool TypeManager::RemoveType(const std::string& type) {
  return removeType(types, type) || removeType(structs, type) ||
         removeType(functions, type);
}

static std::string getKind(const StructUnion& su) {
  return su.isunion ? "union" : "struct";
}

static std::string getKind(const Type& t) { return "typedef"; }

static std::string getKind(const Function& f) { return "function"; }

template <typename K, typename V>
static void enumType(const std::unordered_map<K, V>& map,
                     std::vector<TypeManager::Summary>& types) {
  for (auto i = map.begin(); i != map.end(); ++i) {
    TypeManager::Summary s;
    s.kind = getKind(i->second);
    s.name = i->second.name;
    s.owner = i->second.owner;
    s.size = SizeofType(s.name);
    types.push_back(s);
  }
}

void TypeManager::Enum(std::vector<Summary>& typeList) const {
  typeList.clear();
  enumType(types, typeList);
  enumType(structs, typeList);
  enumType(functions, typeList);
  // nasty hacks to sort in a nice way
  std::sort(typeList.begin(), typeList.end(),
            [](const Summary& a, const Summary& b) {
              auto kindInt = [](const std::string& kind) {
                if (kind == "typedef") return 0;
                if (kind == "struct") return 1;
                if (kind == "union") return 2;
                if (kind == "function") return 3;
                __debugbreak();
                return 4;
              };
              if (a.owner < b.owner)
                return true;
              else if (a.owner > b.owner)
                return false;
              auto ka = kindInt(a.kind), kb = kindInt(b.kind);
              if (ka < kb)
                return true;
              else if (ka > kb)
                return false;
              if (a.name < b.name)
                return true;
              else if (a.name > b.name)
                return false;
              return a.size < b.size;
            });
}

template <typename K, typename V>
static bool mapContains(const std::unordered_map<K, V>& map, const K& k) {
  return map.find(k) != map.end();
}

bool TypeManager::isDefined(const std::string& id) const {
  return mapContains(types, id) || mapContains(structs, id);
}

bool TypeManager::validPtr(const std::string& id) {
  if (id[id.length() - 1] == '*') {
    auto type = id.substr(0, id.length() - 1);
    if (!isDefined(type) && !validPtr(type)) return false;
    std::string owner("ptr");
    auto foundT = types.find(type);
    if (foundT != types.end()) owner = foundT->second.owner;
    auto foundS = structs.find(type);
    if (foundS != structs.end()) owner = foundS->second.owner;
    return addType(owner, Pointer, id, type);
  }
  return false;
}

bool TypeManager::addStructUnion(const StructUnion& s) {
  laststruct = s.name;
  if (s.owner.empty() || s.name.empty() || isDefined(s.name)) return false;
  structs.insert({s.name, s});
  return true;
}

bool TypeManager::addType(const Type& t) {
  if (t.name.empty() || isDefined(t.name)) return false;
  types.insert({t.name, t});
  return true;
}

bool TypeManager::addType(const std::string& owner, Primitive primitive,
                          const std::string& name, const std::string& pointto) {
  if (name.empty() || isDefined(name)) return false;
  Type t;
  t.owner = owner;
  t.name = name;
  t.primitive = primitive;
  t.size = primitivesizes[primitive];
  t.pointto = pointto;
  return addType(t);
}

bool TypeManager::visitMember(const Member& root, Visitor& visitor) const {
  auto foundT = types.find(root.type);
  if (foundT != types.end()) {
    const auto& t = foundT->second;
    if (!t.pointto.empty()) {
      if (!isDefined(t.pointto)) return false;
      if (visitor.visitPtr(root, t))  // allow the visitor to bail out
      {
        if (!Visit(t.pointto, "*" + root.name, visitor)) return false;
        return visitor.visitBack(root);
      }
      return true;
    }
    return visitor.visitType(root, t);
  }
  auto foundS = structs.find(root.type);
  if (foundS != structs.end()) {
    const auto& s = foundS->second;
    if (!visitor.visitStructUnion(root, s)) return false;
    for (const auto& child : s.members) {
      if (child.arrsize) {
        if (!visitor.visitArray(child)) return false;
        for (auto i = 0; i < child.arrsize; i++)
          if (!visitMember(child, visitor)) return false;
        if (!visitor.visitBack(child)) return false;
      } else if (!visitMember(child, visitor))
        return false;
    }
    return visitor.visitBack(root);
  }
  return false;
}

bool AddType(const std::string& owner, const std::string& type,
             const std::string& name) {
  EXCLUSIVE_ACQUIRE(LockTypeManager);
  return typeManager.AddType(owner, type, name);
}

bool AddStruct(const std::string& owner, const std::string& name) {
  EXCLUSIVE_ACQUIRE(LockTypeManager);
  return typeManager.AddStruct(owner, name);
}

bool AddUnion(const std::string& owner, const std::string& name) {
  EXCLUSIVE_ACQUIRE(LockTypeManager);
  return typeManager.AddUnion(owner, name);
}

bool AddMember(const std::string& parent, const std::string& type,
               const std::string& name, int arrsize, int offset) {
  EXCLUSIVE_ACQUIRE(LockTypeManager);
  return typeManager.AddMember(parent, type, name, arrsize, offset);
}

bool AppendMember(const std::string& type, const std::string& name, int arrsize,
                  int offset) {
  EXCLUSIVE_ACQUIRE(LockTypeManager);
  return typeManager.AppendMember(type, name, arrsize, offset);
}

bool AddFunction(const std::string& owner, const std::string& name,
                 const std::string& rettype, Types::CallingConvention callconv,
                 bool noreturn) {
  EXCLUSIVE_ACQUIRE(LockTypeManager);
  return typeManager.AddFunction(owner, name, rettype, callconv, noreturn);
}

bool AddArg(const std::string& function, const std::string& type,
            const std::string& name) {
  EXCLUSIVE_ACQUIRE(LockTypeManager);
  return typeManager.AddArg(function, type, name);
}

bool AppendArg(const std::string& type, const std::string& name) {
  EXCLUSIVE_ACQUIRE(LockTypeManager);
  return typeManager.AppendArg(type, name);
}

int SizeofType(const std::string& type) {
  SHARED_ACQUIRE(LockTypeManager);
  return typeManager.Sizeof(type);
}

bool VisitType(const std::string& type, const std::string& name,
               Types::TypeManager::Visitor& visitor) {
  SHARED_ACQUIRE(LockTypeManager);
  return typeManager.Visit(type, name, visitor);
}

void ClearTypes(const std::string& owner) {
  EXCLUSIVE_ACQUIRE(LockTypeManager);
  return typeManager.Clear(owner);
}

bool RemoveType(const std::string& type) {
  EXCLUSIVE_ACQUIRE(LockTypeManager);
  return typeManager.RemoveType(type);
}

void EnumTypes(std::vector<Types::TypeManager::Summary>& typeList) {
  SHARED_ACQUIRE(LockTypeManager);
  return typeManager.Enum(typeList);
}

int json_default_int(const JSON object, const char* key, int defaultVal) {
  auto jint = json_object_get(object, key);
  if (jint && json_is_integer(jint)) return int(json_integer_value(jint));
  return defaultVal;
}

static void loadTypes(const JSON troot, std::vector<Member>& types) {
  if (!troot) return;
  size_t i;
  JSON vali;
  Member curType;
  json_array_foreach(troot, i, vali) {
    auto type = json_string_value(json_object_get(vali, "type"));
    auto name = json_string_value(json_object_get(vali, "name"));
    if (!type || !*type || !name || !*name) continue;
    curType.type = type;
    curType.name = name;
    types.push_back(curType);
  }
}

static void loadStructUnions(const JSON suroot, bool isunion,
                             std::vector<StructUnion>& structUnions) {
  if (!suroot) return;
  size_t i;
  JSON vali;
  StructUnion curSu;
  curSu.isunion = isunion;
  json_array_foreach(suroot, i, vali) {
    auto suname = json_string_value(json_object_get(vali, "name"));
    if (!suname || !*suname) continue;
    curSu.name = suname;
    curSu.members.clear();
    auto members = json_object_get(vali, "members");
    size_t j;
    JSON valj;
    Member curMember;
    json_array_foreach(members, j, valj) {
      auto type = json_string_value(json_object_get(valj, "type"));
      auto name = json_string_value(json_object_get(valj, "name"));
      if (!type || !*type || !name || !*name) continue;
      curMember.type = type;
      curMember.name = name;
      curMember.arrsize = json_default_int(valj, "arrsize", 0);
      curMember.offset = json_default_int(valj, "offset", -1);
      curSu.members.push_back(curMember);
    }
    structUnions.push_back(curSu);
  }
}

static void loadFunctions(const JSON froot, std::vector<Function>& functions) {
  if (!froot) return;
  size_t i;
  JSON vali;
  Function curFunction;
  json_array_foreach(froot, i, vali) {
    auto rettype = json_string_value(json_object_get(vali, "rettype"));
    auto fname = json_string_value(json_object_get(vali, "name"));
    if (!rettype || !*rettype || !fname || !*fname) continue;
    curFunction.rettype = rettype;
    curFunction.name = fname;
    curFunction.args.clear();
    auto callconv = json_string_value(json_object_get(vali, "callconv"));
    curFunction.noreturn =
        json_boolean_value(json_object_get(vali, "noreturn"));
    if (scmp(callconv, "cdecl"))
      curFunction.callconv = Cdecl;
    else if (scmp(callconv, "stdcall"))
      curFunction.callconv = Stdcall;
    else if (scmp(callconv, "thiscall"))
      curFunction.callconv = Thiscall;
    else if (scmp(callconv, "delphi"))
      curFunction.callconv = Delphi;
    else
      curFunction.callconv = Cdecl;
    auto args = json_object_get(vali, "args");
    size_t j;
    JSON valj;
    Member curArg;
    json_array_foreach(args, j, valj) {
      auto type = json_string_value(json_object_get(valj, "type"));
      auto name = json_string_value(json_object_get(valj, "name"));
      if (!type || !*type || !name || !*name) continue;
      curArg.type = type;
      curArg.name = name;
      curFunction.args.push_back(curArg);
    }
    functions.push_back(curFunction);
  }
}

void LoadModel(const std::string& owner, Model& model) {
  // Add all base struct/union types first to avoid errors later
  for (auto& su : model.structUnions) {
    auto success = su.isunion ? typeManager.AddUnion(owner, su.name)
                              : typeManager.AddStruct(owner, su.name);
    if (!success) {
      // TODO properly handle errors
      dprintf(QT_TRANSLATE_NOOP("DBG", "Failed to add %s %s;\n"),
              su.isunion ? "union" : "struct", su.name.c_str());
      su.name.clear();  // signal error
    }
  }

  // Add simple typedefs
  for (auto& type : model.types) {
    auto success = typeManager.AddType(owner, type.type, type.name);
    if (!success) {
      // TODO properly handle errors
      dprintf(QT_TRANSLATE_NOOP("DBG", "Failed to add typedef %s %s;\n"),
              type.type.c_str(), type.name.c_str());
    }
  }

  // Add base function types to avoid errors later
  for (auto& function : model.functions) {
    auto success =
        typeManager.AddFunction(owner, function.name, function.rettype,
                                function.callconv, function.noreturn);
    if (!success) {
      // TODO properly handle errors
      dprintf(QT_TRANSLATE_NOOP("DBG", "Failed to add function %s %s()\n"),
              function.rettype.c_str(), function.name.c_str());
      function.name.clear();  // signal error
    }
  }

  // Add struct/union members
  for (auto& su : model.structUnions) {
    if (su.name.empty())  // skip error-signalled structs/unions
      continue;
    for (auto& member : su.members) {
      auto success = typeManager.AddMember(su.name, member.type, member.name,
                                           member.arrsize, member.offset);
      if (!success) {
        // TODO properly handle errors
        dprintf(QT_TRANSLATE_NOOP("DBG", "Failed to add member %s %s.%s;\n"),
                member.type.c_str(), su.name.c_str(), member.name.c_str());
      }
    }
  }

  // Add function arguments
  for (auto& function : model.functions) {
    if (function.name.empty())  // skip error-signalled functions
      continue;
    for (auto& arg : function.args) {
      auto success = typeManager.AddArg(function.name, arg.type, arg.name);
      if (!success) {
        // TODO properly handle errors
        dprintf(QT_TRANSLATE_NOOP("DBG", "Failed to add argument %s %s.%s;\n"),
                arg.type.c_str(), function.name.c_str(), arg.name.c_str());
      }
    }
  }
}

bool LoadTypesJson(const std::string& json, const std::string& owner) {
  EXCLUSIVE_ACQUIRE(LockTypeManager);
  auto root = json_loads(json.c_str(), 0, 0);
  if (root) {
    Model model;
    loadTypes(json_object_get(root, "types"), model.types);
    loadTypes(json_object_get(root, ArchValue("types32", "types64")),
              model.types);
    loadStructUnions(json_object_get(root, "structs"), false,
                     model.structUnions);
    loadStructUnions(json_object_get(root, ArchValue("structs32", "structs64")),
                     false, model.structUnions);
    loadStructUnions(json_object_get(root, "unions"), true, model.structUnions);
    loadStructUnions(json_object_get(root, ArchValue("unions32", "unions64")),
                     true, model.structUnions);
    loadFunctions(json_object_get(root, "functions"), model.functions);
    loadFunctions(
        json_object_get(root, ArchValue("functions32", "functions64")),
        model.functions);

    LoadModel(owner, model);

    // Free root
    json_decref(root);
  } else
    return false;
  return true;
}

bool LoadTypesFile(const std::string& path, const std::string& owner) {
  std::string json;
  if (!FileHelper::ReadAllText(path, json)) return false;
  return LoadTypesJson(json, owner);
}
